// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/sparksql/tests/SparkFunctionBaseTest.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions::sparksql::test {
namespace {

class ArrayFlattenTest : public SparkFunctionBaseTest {
 protected:
  void testExpression(
      const std::string& expression,
      const std::vector<VectorPtr>& input,
      const VectorPtr& expected) {
    const auto result = evaluate(expression, make_row_vector(input));
    assertEqualVectors(expected, result);
  }
};

// Flatten integer arrays.
TEST_F(ArrayFlattenTest, intArrays) {
  const auto arrayOfArrays = make_nested_array_vector_from_json<int64_t>({
      "[[1, 1], [2, 2], [3, 3]]",
      "[[4, 4]]",
      "[[5, 5], [6, 6]]",
  });

  const auto expected = make_array_vector_from_json<int64_t>(
      {"[1, 1, 2, 2, 3, 3]", "[4, 4]", "[5, 5, 6, 6]"});

  testExpression("flatten(c0)", {arrayOfArrays}, expected);
}

// Flatten arrays with null.
TEST_F(ArrayFlattenTest, nullArray) {
  const auto arrayOfArrays = make_nested_array_vector_from_json<int64_t>({
      "[[1, 1], null, [3, 3]]",
      "null",
      "[[5, null], [null, 6], [null, null], []]",
  });

  const auto expected = make_array_vector_from_json<int64_t>(
      {"null", "null", "[5, null, null, 6, null, null]"});

  testExpression("flatten(c0)", {arrayOfArrays}, expected);
}
} // namespace
} // namespace kumo::pollux::functions::sparksql::test
